from multiprocessing import (Lock, Event)
from sage.combinat.backtrack import SearchForest
from sage.structure.unique_representation import UniqueRepresentation
from sage.misc.abstract_method import abstract_method
from sage.categories.finite_enumerated_sets import FiniteEnumeratedSets
import multiprocessing
import Queue
import time
import threading
import copy


class SearchForest2(SearchForest):
    r''' TODO '''

    def __init__(self, roots = None, 
                 children = None, 
                 post_process = None,
                 statistic = None, 
                 reduce_function = None, 
                 reduce_init = None, 
                 algorithm = 'depth', 
                 category=None) :
        SearchForest.__init__(self, roots, children, post_process, algorithm, category)
        if post_process is None :
            self.post_process = lambda x: x
        if statistic is not None :
            self.statistic = statistic
        if reduce_function is not None :
            self.reduce_function = reduce_function
        self._res = reduce_init
        self._init = copy.copy(reduce_init)

    @abstract_method
    def statistic(self, o):
        r''' TODO '''
    
    @abstract_method
    def reduce_function(self):
        r''' TODO '''
    
    def run(self):
        r''' 
        TESTS::

            sage: from sage.combinat.Sf import SearchForest2
            sage: S = SearchForest2( [[]],                           # root
            ...     lambda l: [l+[0], l+[1]] if len(l) < 20 else [], # children
            ...     lambda x: x,                                     # post-process
            ...     lambda x: var('x')**len(x),                      # statistic
            ...     lambda x,y: x+y,                                 # reduce_function
            ...     0 )                                              # reduce_init
            sage: S.run()
            1048576*x^20 + 524288*x^19 + 262144*x^18 + 131072*x^17 + 65536*x^16 + 32768*x^15 + 16384*x^14 + 8192*x^13 + 4096*x^12 + 2048*x^11 + 1024*x^10 + 512*x^9 + 256*x^8 + 128*x^7 + 64*x^6 + 32*x^5 + 16*x^4 + 8*x^3 + 4*x^2 + 2*x + 1
        '''
        self._res = copy.copy(self._init)
        self.lock = Lock()
        nbCalculator = max (multiprocessing.cpu_count(), 1)
        event = Event()
        self._thds = [Travelator(self, event) for i in range(nbCalculator)]    # création des threads de parcours
        self._thdsInac = []                                   # liste des threads inactifs

        self._thds[0]._liste.append( iter(self.roots()) )        # le thread 0 reçoit du travail
        for t in self._thds :                                 # on démarre les threads
            t.start()

        event.wait()                                          # attend que les threads dorment

        for t in self._thds :                                 # assemblage des "solutions"
            self._res = self.reduce_function(self._res, t._res)
            t._run = False
            t._liste.append(None)

        return self._res

class Travelator(threading.Thread) :
    """ Classe héritant de thread 
        Cette classe définie une méthode qui parcours l'arbre décrit par SearchForest2.
        A la fin de l'exécution, la variable d'instance _res contient une partie du résultat 
        à combiner avec les résultats des autres threads Travelator.
    """
    def __init__(self, sf2, event) :
        threading.Thread.__init__(self)
        self._obj = sf2
        self.plock = Lock()
        self._event = event
        self._res = copy.copy(self._obj._init)
        self._liste = []
        self._run = True

    def run(self) : 
        """ Corps des threads :
            @note : si on utilise l'event comme indice de boucle??

            @algo :
                tant qu'il y a du travail :
                    on récupère un élément e si possible
                    :: si ce n'est pas possible on cherche un éléments chez les threads frères
                    :: tant qu'on a pas un éléments à traiter on se notifie comme dormeur

                    traitement de e ==> post_process + reduce_function ==> maj de _res
        """
        obj = self._obj
        while self._run :
            nodebf = None
            self.plock.acquire()
            if len(self._liste) > 0 :
                nodebf = self._liste.pop()
                self.plock.release()
            else :
                self.plock.release()
                # la liste est vide
                obj._thdsInac.append(self)
                if len(obj._thds) == len(obj._thdsInac) :
                    # si tout le monde dort, on notifie le patron
                    self._event.set()
                    self._run = False
                    break
                else :
                    # il reste un petit nombre de travailleurs
                    b = True
                    while b and self._run :
                        # on boucle tant qu'il n'y pas de boulot
                        for th in obj._thds :
                            with th.plock :
                                if len(th._liste) > 0 :
                                    nodebf = th._liste.pop()
                                    b = False
                                    break
                    # on se retire de la liste des glandeurs
                    obj._thdsInac.remove(self)
            if nodebf is None :
                continue

            for n in nodebf :
                # pour chaque noeud à traité
                # - maj du resutlat
                self._update(n)
                # - ajout des fils dans la liste
                self._liste.append( iter(obj.children(n)) )

    def _update(self, n) :
        ''' S'occupe de toute la partie maj des valeurs ''reduce_function'' 
        '''
        obj = self._obj
        node = obj.post_process(n)      
        if node is None :
            return

        # maj des valeurs
        self._res = obj.reduce_function(self._res, obj.statistic(node))
